#!/usr/bin/env python
# -*- coding: utf-8 -*-

import os
import argparse
import numpy as np
from tqdm import tqdm
import h5py
import random
from pyntcloud import PyntCloud
import pandas as pd
import urllib.request
import zipfile
import tarfile
import glob
from scipy.spatial import Delaunay
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D

class PointCloudDatasetPreparation:
    def __init__(self, root_dir='./data'):
        """
        初始化点云数据集准备类
        
        Args:
            root_dir: 数据集保存的根目录
        """
        self.root_dir = root_dir
        os.makedirs(root_dir, exist_ok=True)
        
        # 创建子目录
        self.shapenet_dir = os.path.join(root_dir, 'ShapeNet')
        self.modelnet_dir = os.path.join(root_dir, 'ModelNet40')
        self.s3dis_dir = os.path.join(root_dir, 'S3DIS')
        self.kitti_dir = os.path.join(root_dir, 'KITTI')
        self.semantic3d_dir = os.path.join(root_dir, 'Semantic3D')
        self.scannet_dir = os.path.join(root_dir, 'ScanNet')
        
        # 创建噪声数据集目录
        self.noisy_dir = os.path.join(root_dir, 'noisy_datasets')
        os.makedirs(self.noisy_dir, exist_ok=True)
        
    def download_shapenet(self):
        """
        下载ShapeNet数据集（合成3D模型）
        注意：需要申请权限才能下载完整ShapeNet，这里使用一部分示例
        """
        print("注意：完整ShapeNet数据集需要在官方网站注册并授权下载")
        print("下载ShapeNetCore简化版本...")
        
        os.makedirs(self.shapenet_dir, exist_ok=True)
        # ShapeNet数据集下载URL（以v2版本为例）
        url = "https://shapenet.cs.stanford.edu/shapenet/obj-zip/ShapeNetCore.v2.zip"
        
        # 实际项目中请去官方网站申请完整数据集：https://shapenet.org/
        print("请从ShapeNet官方网站申请完整数据集")
        print("示例数据使用ShapeNetCore.v2中的部分类别")
        
        # 此处仅作为示例，实际执行时需要根据授权下载
    
    def download_modelnet40(self):
        """
        下载ModelNet40数据集（合成3D模型）
        """
        print("下载ModelNet40数据集...")
        
        os.makedirs(self.modelnet_dir, exist_ok=True)
        url = "https://shapenet.cs.stanford.edu/media/modelnet40_ply_hdf5_2048.zip"
        zip_path = os.path.join(self.modelnet_dir, "modelnet40_ply_hdf5_2048.zip")
        
        # 下载文件
        if not os.path.exists(zip_path):
            print(f"从 {url} 下载ModelNet40数据集...")
            urllib.request.urlretrieve(url, zip_path)
            
            # 解压文件
            with zipfile.ZipFile(zip_path, 'r') as zip_ref:
                zip_ref.extractall(self.modelnet_dir)
            
            print("ModelNet40数据集下载并解压完成")
        else:
            print("ModelNet40数据集已存在")
    
    def download_s3dis(self):
        """
        下载S3DIS数据集（真实室内扫描）
        注意：S3DIS数据集需要在官方网站注册
        """
        print("注意：S3DIS数据集需要在官方网站注册才能下载")
        print("请访问：http://buildingparser.stanford.edu/dataset.html")
        
        os.makedirs(self.s3dis_dir, exist_ok=True)
    
    def download_kitti(self):
        """
        下载KITTI数据集（自动驾驶场景的点云数据）
        """
        print("下载KITTI数据集的点云部分...")
        
        os.makedirs(self.kitti_dir, exist_ok=True)
        # KITTI数据集下载URL（以点云部分为例）
        url = "https://s3.eu-central-1.amazonaws.com/avg-kitti/data_object_velodyne.zip"
        
        print("注意：KITTI数据集需要在官方网站注册才能下载完整数据")
        print("请访问：http://www.cvlibs.net/datasets/kitti/")
        
        # 此处仅作为示例，实际执行时根据需要下载
    
    def download_semantic3d(self):
        """
        下载Semantic3D数据集（大规模室外扫描）
        """
        print("下载Semantic3D数据集的示例部分...")
        
        os.makedirs(self.semantic3d_dir, exist_ok=True)
        # 完整数据集很大，这里只给出下载链接
        print("请从官方网站下载完整数据：http://www.semantic3d.net/")
        # 可以选择下载reduced-8数据集，规模更小
    
    def download_scannet(self):
        """
        下载ScanNet数据集（室内场景扫描）
        """
        print("ScanNet数据集需要申请权限才能下载")
        print("请访问：http://www.scan-net.org/")
        
        os.makedirs(self.scannet_dir, exist_ok=True)
    
    def generate_synthetic_data(self, num_points=1024, num_shapes=100):
        """
        生成合成的点云数据
        
        Args:
            num_points: 每个点云的点数
            num_shapes: 生成的形状数量
        """
        print(f"生成{num_shapes}个合成点云形状...")
        
        synthetic_dir = os.path.join(self.root_dir, 'synthetic')
        os.makedirs(synthetic_dir, exist_ok=True)
        
        shapes = []
        
        # 生成基本几何形状
        for i in range(num_shapes):
            shape_type = random.choice(['sphere', 'cube', 'cylinder', 'plane'])
            
            if shape_type == 'sphere':
                # 生成球体
                points = self._generate_sphere(num_points, radius=random.uniform(0.5, 1.5))
                
            elif shape_type == 'cube':
                # 生成立方体
                width = random.uniform(0.5, 2.0)
                height = random.uniform(0.5, 2.0)
                depth = random.uniform(0.5, 2.0)
                points = self._generate_cube(num_points, width, height, depth)
                
            elif shape_type == 'cylinder':
                # 生成圆柱体
                radius = random.uniform(0.5, 1.5)
                height = random.uniform(1.0, 3.0)
                points = self._generate_cylinder(num_points, radius, height)
                
            else:  # plane
                # 生成平面
                width = random.uniform(1.0, 3.0)
                height = random.uniform(1.0, 3.0)
                points = self._generate_plane(num_points, width, height)
            
            # 随机旋转和平移
            points = self._random_transform(points)
            
            # 保存点云
            shapes.append(points)
            
            # 保存为ply文件
            self._save_point_cloud_ply(
                os.path.join(synthetic_dir, f"shape_{i:03d}_{shape_type}.ply"), 
                points
            )
        
        # 保存为h5文件
        with h5py.File(os.path.join(synthetic_dir, 'synthetic_shapes.h5'), 'w') as f:
            f.create_dataset('points', data=np.array(shapes))
            f.create_dataset('labels', data=np.zeros(len(shapes)))  # 默认标签为0
        
        print(f"已生成{num_shapes}个合成点云形状，保存在 {synthetic_dir}")
        return synthetic_dir
    
    def _generate_sphere(self, num_points, radius=1.0):
        """生成球体点云"""
        points = []
        for _ in range(num_points):
            # 在球体上均匀采样点
            phi = np.random.uniform(0, 2 * np.pi)
            costheta = np.random.uniform(-1, 1)
            theta = np.arccos(costheta)
            
            x = radius * np.sin(theta) * np.cos(phi)
            y = radius * np.sin(theta) * np.sin(phi)
            z = radius * np.cos(theta)
            
            points.append([x, y, z])
        
        return np.array(points)
    
    def _generate_cube(self, num_points, width, height, depth):
        """生成立方体点云"""
        points = []
        
        # 均匀分配到6个面
        points_per_face = num_points // 6
        
        # 生成每个面的点
        for face in range(6):
            for _ in range(points_per_face):
                if face == 0:  # 正x面
                    x = width/2
                    y = np.random.uniform(-height/2, height/2)
                    z = np.random.uniform(-depth/2, depth/2)
                elif face == 1:  # 负x面
                    x = -width/2
                    y = np.random.uniform(-height/2, height/2)
                    z = np.random.uniform(-depth/2, depth/2)
                elif face == 2:  # 正y面
                    x = np.random.uniform(-width/2, width/2)
                    y = height/2
                    z = np.random.uniform(-depth/2, depth/2)
                elif face == 3:  # 负y面
                    x = np.random.uniform(-width/2, width/2)
                    y = -height/2
                    z = np.random.uniform(-depth/2, depth/2)
                elif face == 4:  # 正z面
                    x = np.random.uniform(-width/2, width/2)
                    y = np.random.uniform(-height/2, height/2)
                    z = depth/2
                else:  # 负z面
                    x = np.random.uniform(-width/2, width/2)
                    y = np.random.uniform(-height/2, height/2)
                    z = -depth/2
                
                points.append([x, y, z])
        
        # 添加剩余的点
        remaining = num_points - len(points)
        for _ in range(remaining):
            face = np.random.randint(0, 6)
            if face == 0:  # 正x面
                x = width/2
                y = np.random.uniform(-height/2, height/2)
                z = np.random.uniform(-depth/2, depth/2)
            elif face == 1:  # 负x面
                x = -width/2
                y = np.random.uniform(-height/2, height/2)
                z = np.random.uniform(-depth/2, depth/2)
            elif face == 2:  # 正y面
                x = np.random.uniform(-width/2, width/2)
                y = height/2
                z = np.random.uniform(-depth/2, depth/2)
            elif face == 3:  # 负y面
                x = np.random.uniform(-width/2, width/2)
                y = -height/2
                z = np.random.uniform(-depth/2, depth/2)
            elif face == 4:  # 正z面
                x = np.random.uniform(-width/2, width/2)
                y = np.random.uniform(-height/2, height/2)
                z = depth/2
            else:  # 负z面
                x = np.random.uniform(-width/2, width/2)
                y = np.random.uniform(-height/2, height/2)
                z = -depth/2
            
            points.append([x, y, z])
        
        return np.array(points)
    
    def _generate_cylinder(self, num_points, radius, height):
        """生成圆柱体点云"""
        points = []
        
        # 90% 的点在侧面，5% 在顶部，5% 在底部
        side_points = int(num_points * 0.9)
        cap_points = int(num_points * 0.05)
        
        # 侧面点
        for _ in range(side_points):
            theta = np.random.uniform(0, 2 * np.pi)
            h = np.random.uniform(-height/2, height/2)
            
            x = radius * np.cos(theta)
            y = radius * np.sin(theta)
            z = h
            
            points.append([x, y, z])
        
        # 顶部和底部点
        for _ in range(cap_points):
            r = radius * np.sqrt(np.random.uniform(0, 1))
            theta = np.random.uniform(0, 2 * np.pi)
            
            x = r * np.cos(theta)
            y = r * np.sin(theta)
            z = height/2
            
            points.append([x, y, z])
        
        for _ in range(cap_points):
            r = radius * np.sqrt(np.random.uniform(0, 1))
            theta = np.random.uniform(0, 2 * np.pi)
            
            x = r * np.cos(theta)
            y = r * np.sin(theta)
            z = -height/2
            
            points.append([x, y, z])
        
        # 添加剩余的点
        remaining = num_points - len(points)
        for _ in range(remaining):
            theta = np.random.uniform(0, 2 * np.pi)
            h = np.random.uniform(-height/2, height/2)
            
            x = radius * np.cos(theta)
            y = radius * np.sin(theta)
            z = h
            
            points.append([x, y, z])
        
        return np.array(points)
    
    def _generate_plane(self, num_points, width, height):
        """生成平面点云"""
        plane_points = []
        for _ in range(num_points):
            x = np.random.uniform(-width/2, width/2)
            y = np.random.uniform(-height/2, height/2)
            z = np.random.uniform(-0.05, 0.05)  # 添加微小的z轴变化
            plane_points.append([x, y, z])
        
        return np.array(plane_points)
    
    def _random_transform(self, points):
        """随机旋转和平移点云"""
        # 随机旋转矩阵
        theta_x = np.random.uniform(0, 2*np.pi)
        theta_y = np.random.uniform(0, 2*np.pi)
        theta_z = np.random.uniform(0, 2*np.pi)
        
        Rx = np.array([
            [1, 0, 0],
            [0, np.cos(theta_x), -np.sin(theta_x)],
            [0, np.sin(theta_x), np.cos(theta_x)]
        ])
        
        Ry = np.array([
            [np.cos(theta_y), 0, np.sin(theta_y)],
            [0, 1, 0],
            [-np.sin(theta_y), 0, np.cos(theta_y)]
        ])
        
        Rz = np.array([
            [np.cos(theta_z), -np.sin(theta_z), 0],
            [np.sin(theta_z), np.cos(theta_z), 0],
            [0, 0, 1]
        ])
        
        R = np.dot(Rz, np.dot(Ry, Rx))
        
        # 随机平移向量
        t = np.array([
            np.random.uniform(-2, 2),
            np.random.uniform(-2, 2),
            np.random.uniform(-2, 2)
        ])
        
        # 应用变换
        transformed_points = np.dot(points, R.T) + t
        
        return transformed_points
    
    def _save_point_cloud_ply(self, filename, points):
        """保存点云为ply文件"""
        df = pd.DataFrame(
            data=points,
            columns=['x', 'y', 'z']
        )
        cloud = PyntCloud(df)
        cloud.to_file(filename)
    
    def _load_point_cloud(self, filename):
        """读取点云文件"""
        cloud = PyntCloud.from_file(filename)
        points = cloud.points[['x', 'y', 'z']].values
        return points
    
    def add_noise(self, dataset_path, noise_types=['gaussian', 'outlier', 'density'], noise_levels=[0.01, 0.02, 0.05]):
        """
        为点云数据集添加不同类型和级别的噪声
        
        Args:
            dataset_path: 原始数据集路径
            noise_types: 噪声类型列表
            noise_levels: 噪声级别列表
        """
        print(f"为数据集 {dataset_path} 添加噪声...")
        
        # 获取目录中的所有点云文件
        point_cloud_files = []
        for ext in ['*.ply', '*.pcd', '*.xyz']:
            point_cloud_files.extend(glob.glob(os.path.join(dataset_path, ext)))
        
        # 如果是h5文件，则直接处理h5文件
        h5_files = glob.glob(os.path.join(dataset_path, '*.h5'))
        
        if not point_cloud_files and not h5_files:
            print(f"在 {dataset_path} 中未找到点云文件")
            return
        
        # 为每种噪声类型和级别创建目录
        for noise_type in noise_types:
            for noise_level in noise_levels:
                noisy_subdir = os.path.join(
                    self.noisy_dir,
                    f"{os.path.basename(dataset_path)}_{noise_type}_{noise_level}"
                )
                os.makedirs(noisy_subdir, exist_ok=True)
                
                # 处理单个点云文件
                for pc_file in tqdm(point_cloud_files, desc=f"添加{noise_type}噪声(级别:{noise_level})"):
                    try:
                        # 读取点云
                        points = self._load_point_cloud(pc_file)
                        
                        # 添加噪声
                        noisy_points = self._add_noise_to_points(points, noise_type, noise_level)
                        
                        # 保存带有噪声的点云
                        output_file = os.path.join(noisy_subdir, os.path.basename(pc_file))
                        self._save_point_cloud_ply(output_file, noisy_points)
                    except Exception as e:
                        print(f"处理文件 {pc_file} 时出错: {e}")
                
                # 处理h5文件
                for h5_file in h5_files:
                    try:
                        output_h5 = os.path.join(
                            noisy_subdir,
                            f"noisy_{os.path.basename(h5_file)}"
                        )
                        
                        # 读取h5文件
                        with h5py.File(h5_file, 'r') as f_in, h5py.File(output_h5, 'w') as f_out:
                            # 处理points数据集
                            if 'points' in f_in:
                                points = f_in['points'][:]
                                
                                # 为每个点云添加噪声
                                noisy_points = np.zeros_like(points)
                                for i in tqdm(range(len(points)), desc=f"处理h5文件中的点云"):
                                    noisy_points[i] = self._add_noise_to_points(points[i], noise_type, noise_level)
                                
                                # 保存带有噪声的点云
                                f_out.create_dataset('points', data=noisy_points)
                                
                                # 复制其他数据集
                                for key in f_in.keys():
                                    if key != 'points':
                                        f_out.create_dataset(key, data=f_in[key][:])
                    except Exception as e:
                        print(f"处理h5文件 {h5_file} 时出错: {e}")
        
        print(f"噪声添加完成，结果保存在 {self.noisy_dir}")
    
    def _add_noise_to_points(self, points, noise_type, noise_level):
        """
        为点云添加指定类型和级别的噪声
        
        Args:
            points: 点云数据 (N, 3)
            noise_type: 噪声类型
            noise_level: 噪声级别
            
        Returns:
            带有噪声的点云
        """
        noisy_points = points.copy()
        
        if noise_type == 'gaussian':
            # 添加高斯噪声
            noise = np.random.normal(0, noise_level, size=points.shape)
            noisy_points = points + noise
            
        elif noise_type == 'outlier':
            # 添加离群点
            num_outliers = int(len(points) * noise_level)
            
            # 计算点云的边界盒
            min_bound = np.min(points, axis=0)
            max_bound = np.max(points, axis=0)
            
            # 生成随机的离群点
            outliers = np.random.uniform(
                min_bound - (max_bound - min_bound) * 0.2,
                max_bound + (max_bound - min_bound) * 0.2,
                size=(num_outliers, 3)
            )
            
            # 随机选择要替换的点的索引
            indices = np.random.choice(len(points), num_outliers, replace=False)
            noisy_points[indices] = outliers
            
        elif noise_type == 'density':
            # 模拟密度不均匀
            # 随机删除一些点
            num_to_keep = int(len(points) * (1 - noise_level))
            indices = np.random.choice(len(points), num_to_keep, replace=False)
            noisy_points = points[indices]
            
            # 如果需要保持点的数量不变，可以复制一些点
            if len(noisy_points) < len(points):
                extra_indices = np.random.choice(len(noisy_points), len(points) - len(noisy_points), replace=True)
                noisy_points = np.vstack([noisy_points, noisy_points[extra_indices]])
        
        else:
            print(f"未知的噪声类型: {noise_type}")
        
        return noisy_points
    
    def visualize_point_cloud(self, pcd_path):
        """
        可视化点云
        
        Args:
            pcd_path: 点云文件路径
        """
        # 读取点云
        cloud = PyntCloud.from_file(pcd_path)
        points = cloud.points[['x', 'y', 'z']].values
        
        # 创建3D图
        fig = plt.figure(figsize=(10, 8))
        ax = fig.add_subplot(111, projection='3d')
        
        # 绘制点云
        ax.scatter(points[:, 0], points[:, 1], points[:, 2], s=1, c=points[:, 2], cmap='viridis')
        
        # 设置坐标轴标签
        ax.set_xlabel('X')
        ax.set_ylabel('Y')
        ax.set_zlabel('Z')
        
        # 设置标题
        ax.set_title(f'点云可视化: {os.path.basename(pcd_path)}')
        
        # 显示图形
        plt.tight_layout()
        plt.show()
    
    def visualize_noise_comparison(self, original_path, noisy_path):
        """
        对比可视化原始点云和带噪声的点云
        
        Args:
            original_path: 原始点云文件路径
            noisy_path: 带噪声的点云文件路径
        """
        # 读取点云
        original_cloud = PyntCloud.from_file(original_path)
        noisy_cloud = PyntCloud.from_file(noisy_path)
        
        original_points = original_cloud.points[['x', 'y', 'z']].values
        noisy_points = noisy_cloud.points[['x', 'y', 'z']].values
        
        # 创建图像
        fig = plt.figure(figsize=(12, 6))
        
        # 原始点云
        ax1 = fig.add_subplot(121, projection='3d')
        ax1.scatter(original_points[:, 0], original_points[:, 1], original_points[:, 2], 
                   s=1, c='blue', alpha=0.8)
        ax1.set_title('原始点云')
        ax1.set_xlabel('X')
        ax1.set_ylabel('Y')
        ax1.set_zlabel('Z')
        
        # 带噪声的点云
        ax2 = fig.add_subplot(122, projection='3d')
        ax2.scatter(noisy_points[:, 0], noisy_points[:, 1], noisy_points[:, 2], 
                   s=1, c='orange', alpha=0.8)
        ax2.set_title('带噪声的点云')
        ax2.set_xlabel('X')
        ax2.set_ylabel('Y')
        ax2.set_zlabel('Z')
        
        # 显示图形
        plt.tight_layout()
        plt.show()
    
    def prepare_datasets(self):
        """
        准备所有数据集
        """
        print("开始准备点云数据集...")
        
        # 下载和处理公开数据集
        self.download_modelnet40()  # 合成数据
        # 以下数据集通常需要注册才能下载，所以这里只提供提示信息
        # self.download_shapenet()
        # self.download_s3dis()
        # self.download_kitti()
        # self.download_semantic3d()
        # self.download_scannet()
        
        # 生成合成数据
        synthetic_dir = self.generate_synthetic_data(num_points=1024, num_shapes=100)
        
        # 添加噪声
        self.add_noise(synthetic_dir)
        
        if os.path.exists(os.path.join(self.modelnet_dir, "modelnet40_ply_hdf5_2048")):
            modelnet_h5_dir = os.path.join(self.modelnet_dir, "modelnet40_ply_hdf5_2048")
            self.add_noise(modelnet_h5_dir)
        
        print("数据集准备完成！")

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description='点云数据集准备工具')
    parser.add_argument('--root_dir', type=str, default='./data', help='数据集保存的根目录')
    parser.add_argument('--action', type=str, default='prepare_all', 
                        choices=['prepare_all', 'download_modelnet', 'generate_synthetic',
                                'add_noise', 'visualize'],
                        help='执行的操作')
    parser.add_argument('--pcd_path', type=str, default=None, help='要可视化的点云文件路径')
    
    args = parser.parse_args()
    
    dataset_prep = PointCloudDatasetPreparation(root_dir=args.root_dir)
    
    if args.action == 'prepare_all':
        dataset_prep.prepare_datasets()
    elif args.action == 'download_modelnet':
        dataset_prep.download_modelnet40()
    elif args.action == 'generate_synthetic':
        dataset_prep.generate_synthetic_data()
    elif args.action == 'add_noise':
        if not args.pcd_path:
            print("请指定点云文件或目录路径")
        else:
            dataset_prep.add_noise(args.pcd_path)
    elif args.action == 'visualize':
        if not args.pcd_path:
            print("请指定点云文件路径")
        else:
            dataset_prep.visualize_point_cloud(args.pcd_path) 